""" Module containing functions for working with GIMIC data.  Read the
function docstrings for more info.

Examples
--------

>>> import gimic
>>> snapdir = '/path/to/GIMIC/snapshot/dir/'
>>> i,z = gimic.find_isnap(snapdir)
>>> masses = gimic.readsnap(i, 3, 'ParticleMass', indir=snapdir)
"""
import os
import glob
import numpy as np
import tables
import pdb

def find_isnap(snapdir, redshift):
    """ For a directory of snapshots, find the snapshot number with
    redshift closest to the requested redshift.

    Inputs
    ------
    snapdir   str       Snapshot directory
    redshift  float     Desired redshift

    Returns
    -------
    isnap, zsnap   (int, float)
        Number and redshift of the snapshot that is as close as
        possible to the input redshift.
    """

    nsnap = len(glob.glob(os.path.join(simdir, 'snapshot_*.hdf5')))
    if nsnap == 0:
        raise ValueError('No snapshots found in %s!' % snapdir)

    zbest = 1.e10

    # read the redshift of every snapshot file, and find the one with
    # redshift closest to the requested redshift.
    for i in reversed(range(nsnap)):
        filename = '%s/snapshot_%03i/snap_%03i.9.hdf5' % (simdir, i, i)
        fh = tables.openFile(filename)
        z = fh.getNodeAttr('/Header','Redshift')
        fh.close()
        if abs(z - redshift) <= abs(zbest - redshift):
            ibest = i
            zbest = z
        else:
            break

    return ibest, zbest


def readsnap(isnap, itype, variable, indir='.', prefixdir='snapshot_', 
             prefixfile='snap_', verbose=False):
    """ Read particle data from a GIMIC snapshot.

    Input
    -----
    isnap: int
        Snapshot number (found with, e.g. find_isnap).
    itype: int
        Particle type number. 0 <= itype <= 5.
    variable: str
        Name of variable to read.
    indir: str ('.')
        Directory holding snapshots. 
    prefixdir: str ('snapshot_')
        Prefix of an individual snapshot directory.
    prefixfile: str ('snap_')
        Prefix of an individual snapshot file.
    verbose: bool (False)
        Print helpful messages.

    Returns
    -------
    out:  ndarray
        Array of particle data.
    """

    # sometimes there's just a single file? Ignore for now since it
    # simplifies the code.

#     nfiles = 1
#     infile = os.path.join(indir, '%s%03i' % (prefixdir, isnap), 
#                           '%s%03i.hdf5' % (prefixfile, isnap))

    infile = os.path.join(indir, '%s%03i' % (prefixdir, isnap),
                          '%s%03i.0.hdf5' % (prefixfile, isnap) )

    fh = tables.openFile(infile)

    # number of files for this snapshot
    nfiles = fh.getNodeAttr('/Header', 'NumFilesPerSnapshot')

    if verbose: 
        # total number of particles in the snapshot
        ntotal = fh.getNodeAttr('/Header', 'NumPart_Total')
        print '%s particles of type %i in total' % (itype, ntotal[itype])

    fh.close()

    # loop over files to collect data
    varname = '/PartType%i/%s' % (itype, variable)
    total = []
    for i in range(nfiles):
        infile = os.path.join(indir, '%s%03i' % (prefixdir, isnap),
                              '%s%03i.%i.hdf5' % (prefixfile, isnap, i))
        fh = tables.openFile(infile)
        if verbose:
            npart = fh.getNodeAttr('/Header', 'NumPart_ThisFile')
            print ('readsnap(): reading file %s - it contains %i particles'
                   % (infile, npart[itype]))
        total.append(fh.getNode(varname).read())
        fh.close()

    total = np.concatenate(total)

    return total


def write_gals(region, resolution, redshift, minmass=0.01, saferadius=15.):
    """ Find galaxies above some mass threshhold for a given region,
    resolution and redshift.

    Inputs
    ------
    resolution: str
        'Lo'. 'Hi' not implemented yet

    region: int
        One of -2, -1, 0, 1, or 2.

    redshift: float

    minmass: float (0.01)
        Only select galaxies with stellar masses above this many times
        10**10 solar masses.

    saferadius: float (15.0)
        Radius of the 'safe' region in h**-1 comoving Mpc.

    Notes
    -----
    Galaxy data is written to a text file gals_*_snap???.txt
    """
    regions = {-2: 'Sigmam2',
               -1: 'Sigmam1',
               0 : 'Sigma0',
               1 : 'Sigmap1',
               2 : 'Sigmap2'}
    reg_str = regions[region]
    res_str = resolution.capitalize() + 'Res'
    prefix = os.path.join('/data/gimic/rcrain/', res_str, reg_str, 'Data/')

    # find the snapshot number corresponding to the requested redshift.
    isnap,zsnap = find_isnap(prefix + 'Snapshots/', redshift)

    # read subhalo masses and positions
    subdir = prefix + 'SubGroups/'
    subtabname = subdir + 'groups_%03i/subhalo_tab_%03i.hdf5' % (
        isnap, isnap)
    print 'Reading subhalo info from', subtabname
    fh = tables.openFile(subtabname)

    # read subhalo properties
    stelmass = fh.root.SubHalo.MassTab.read()[:,4]
    totmass = fh.root.SubHalo.Mass.read()
    x, y, z = fh.root.SubHalo.Position.read().T
    vx, vy, vz = fh.root.SubHalo.Velocity.read().T
    ifof = fh.root.SubHalo.FOFGroupNr.read()
    fh.close()
    names = 'x y z vx vy vz mtot mstel ifof'.split()
    subhalo = np.rec.fromarrays([x,y,z,vx,vy,vz,totmass,stelmass,ifof],
                                names=names)

    # find the region's centre of (dark matter?) mass x,y,z
    # coordinates comoving h^-1 Mpc
    filename = os.path.join('/data/rw12/tt/GIMIC/', reg_str,
                            'Info/Info.snap_%03i.hdf5' % isnap)
    fh = tables.openFile(filename)
    xc, yc, zc = fh.root.PartType1.CMS.read()
    fh.close()

    print 'Selecting galaxies'
    # select indices of good galaxies, good meaning:
    # subhalo stellar mass bigger than 10**10 solar masses

    # make sure stellar matter (particle number 4) is above the
    # minimum mass
    c0 = stelmass >= minmass
    # subhalo mass at least 90% of the mass of FOF halo that contains it
    #c1 = subtab.subhalo.mass GE 0.9*subtab.fof.mass[subtab.subhalo.FoFGroupNr]
    # inside safe region

    subhalo1 = subhalo[c0]

    filename = 'gals_%s%s_snap%03i.txt' % (res_str, reg_str, isnap)
    print 'Printing to ', filename
    fh = open(filename,'w')
    # print a header with gimic simulation info

    s =  ('# subhalos with a minimum stellar mass of %.3f '
          '10**10 solar masses\n' % minmass)
    s += '# For %+i %s gimic redshift=%.5f snapshot\n' % (region,res_str,zsnap)
    s += '# centre position of volume x,y,z = %.2f, %.2f, %.2f\n' % (xc,yc,zc)
    s += '# volume radius %.2f h^-1 Mpc\n' % saferadius
    s += """\
# 1. x position / (h^-1 Mpc), comoving
# 2. y position / (h^-1 Mpc), comoving
# 3. z position / (h^-1 Mpc), comoving
# 4. x velocity / (km/s)
# 5. y velocity / (km/s)
# 6. z velocity / (km/s)
# 7. stellar mass / (10**10 solar masses)
# 8. total mass / (10**10 solar masses)
# 9. ID number of parent FOF halo
"""
    fh.write(s)
    
    for s in subhalo1:
        fh.write(('%.5f '*8 + '%i\n') % tuple(s[k] for k in names))
        
    fh.close()


def cic(value, x, nx, y=None, ny=1, z=None, nz=1,
        wraparound=False, average=True):
    """ Interpolate an irregularly sampled field using Cloud in Cell
    method.

    This function interpolates an irregularly sampled field to a
    regular grid using Cloud In Cell (nearest grid point gets weight
    1-dngp, point on other side gets weight dngp, where dngp is the
    distance to the nearest grid point in units of the cell size).
    
    Inputs
    ------
    value: array, shape (N,)
        Sample weights (field values). For a temperature field this
        would be the temperature and the keyword average should be
        True. For a density field this could be either the particle
        mass (average should be False) or the density (average should
        be True).
    x: array, shape (N,)
        X coordinates of field samples, unit indices: [0,NX>.
    nx: int
        Number of grid points in X-direction.
    y: array, shape (N,), optional
        Y coordinates of field samples, unit indices: [0,NY>.
    ny: int, optional
        Number of grid points in Y-direction.
    z: array, shape (N,), optional
        Z coordinates of field samples, unit indices: [0,NZ>.
    nz: int, optional
        Number of grid points in Z-direction.
    wraparound: bool (False)
        If True, then values past the first or last grid point can
        wrap around and contribute to the grid point on the opposite
        side (see the Notes section below).
    average: bool (False)
        If True, average the contributions of each value to a grid
        point instead of summing them.

    Returns
    -------
    dens: ndarray, shape (nx, ny, nz)
        The grid point values.

    Notes
    -----
    Example of default allocation of nearest grid points: nx = 4, * = gridpoint.

      0   1   2   3     Index of gridpoints
      *   *   *   *     Grid points
    |---|---|---|---|   Range allocated to gridpoints ([0.0,1.0> -> 0, etc.)
    0   1   2   3   4   posx

    Example of ngp allocation for wraparound=True: nx = 4, * = gridpoint.

      0   1   2   3        Index of gridpoints
      *   *   *   *        Grid points
    |---|---|---|---|--    Range allocated to gridpoints ([0.5,1.5> -> 1, etc.)
      0   1   2   3   4=0  posx


    References
    ----------
    R.W. Hockney and J.W. Eastwood, Computer Simulations Using Particles
        (New York: McGraw-Hill, 1981).

    Modification History
    --------------------
    IDL code written by Joop Schaye, Feb 1999.
    Avoid integer overflow for large dimensions P.Riley/W.Landsman Dec. 1999
    Translated to Python by Neil Crighton, July 2009.
    
    Examples
    --------
    >>> nx = 20
    >>> ny = 10
    >>> posx = np.random.rand(size=1000)
    >>> posy = np.random.rand(size=1000)
    >>> value = posx**2 + posy**2
    >>> field = cic(value, posx*nx, nx, posy*ny, ny)
    # plot surface
    """

    def findweights(pos, ngrid):
        """ Calculate CIC weights.
        
        Coordinates of nearest grid point (ngp) to each value. """

        if wraparound:
            # grid points at integer values
            ngp = np.fix(pos + 0.5)
        else:
            # grid points are at half-integer values, starting at 0.5,
            # ending at len(grid) - 0.5
            ngp = np.fix(pos) + 0.5

        # Distance from sample to ngp.
        distngp = ngp - pos

        # weight for higher (right, w2) and lower (left, w1) ngp
        weight2 = np.abs(distngp)
        weight1 = 1.0 - weight2

        # indices of the nearest grid points
        if wraparound:
            ind1 = ngp
        else:
            ind1 = ngp - 0.5
        ind1 = ind1.astype(int)

        ind2 = ind1 - 1
        # Correct points where ngp < pos (ngp to the left).
        ind2[distngp < 0] += 2

        # Note that ind2 can be both -1 and ngrid at this point,
        # regardless of wraparound. This is because distngp can be
        # exactly zero.
        bad = (ind2 == -1)
        ind2[bad] = ngrid - 1
        if not wraparound:
            weight2[bad] = 0.
        bad = (ind2 == ngrid)
        ind2[bad] = 0
        if not wraparound:
            weight2[bad] = 0.

        if wraparound:
            ind1[ind1 == ngrid] = 0

        return dict(weight=weight1, ind=ind1), dict(weight=weight2, ind=ind2)


    def update_field_vals(field, totalweight, a, b, c, value, debug=True):
        """ This updates the field array (and the totweight array if
        average is True).

        The elements to update and their values are inferred from
        a,b,c and value.
        """
        print 'Updating field vals'
        # indices for field - doesn't include all combinations
        indices = a['ind'] + b['ind'] * nx + c['ind'] * nxny
        # weight per coordinate
        weights = a['weight'] * b['weight'] * c['weight']
        # Don't modify the input value array, just rebind the name.
        value = weights * value 
        if average:
            for i,ind in enumerate(indices):
                field[ind] += value[i]
                totalweight[ind] += weights[i]
        else:
            for i,ind in enumerate(indices):
                field[ind] += value[i]
            #if debug: print ind, weights[i], value[i], field[ind]


    #pdb.set_trace()
    nx, ny, nz = (int(i) for i in (nx, ny, nz))
    nxny = nx * ny
    value = np.asarray(value)

    print 'Resampling %i values to a %i by %i by %i grid' % (
        len(value), nx, ny, nz)

    # normalise data such that grid points are at integer positions.
    #x = (x - x.min()) / x.ptp() * nx
    #y = (y - y.min()) / y.ptp() * ny
    #z = (z - z.min()) / z.ptp() * nz

    x1, x2 = findweights(np.asarray(x), nx)
    y1 = z1 = dict(weight=1., ind=0)
    if y is not None:
        y1, y2 = findweights(np.asarray(y), ny)
        if z is not None:
            z1, z2 = findweights(np.asarray(z), nz)

    # float32 to save memory for big arrays (e.g. 256**3)
    field = np.zeros(nx * ny * nz, np.float32)

    if average:
        totalweight = np.zeros(nx * ny * nz, np.float32)
    else:
        totalweight = None

    update_field_vals(field, totalweight, x1, y1, z1, value)
    update_field_vals(field, totalweight, x2, y1, z1, value)
    if y is not None:
        update_field_vals(field, totalweight, x1, y2, z1, value)
        update_field_vals(field, totalweight, x2, y2, z1, value)
        if z is not None:
            update_field_vals(field, totalweight, x1, y1, z2, value)
            update_field_vals(field, totalweight, x2, y1, z2, value)
            update_field_vals(field, totalweight, x1, y2, z2, value)
            update_field_vals(field, totalweight, x2, y2, z2, value)

    if average:
        good = totalweight > 0
        field[good] /= totalweight[good]

    return field.reshape((nx, ny, nz)).squeeze()

def example_cic():
    from astro.gimic import cic
    # 1d
    x = np.random.rand(10000)*10
    vals = x**2
    pl.figure()
    plot(x, vals, '.')
    field = cic(vals, x, 10)
    plot(np.arange(10), field, ls='steps-post')
    # 2d
    y = np.random.rand(10000)*10
    vals = x**2 + y**2
    field = cic(vals, x, 1000, y, 1000)

    # colour represents field value
    pl.figure()
    pl.pcolor(field, vmin=0, vmax=2*10**2)
    pl.scatter(x, y, c=vals, vmin=0, vmax=2*10**2)
    pl.colorbar()
    pl.show()

    # 3d
    z = np.random.rand(10000)*10
    vals = x**2 + y**2 + z**2
    field = cic(vals, x, 200, y, 200, z, 200)

    #from mpl_toolkits.mplot3d import Axes3D
    # fig = pl.figure()
    # ax = Axes3D(fig)
    # X = np.arange(10)
    # Y = np.arange(10)
    # X, Y = np.meshgrid(X, Y)
    # ax.plot_wireframe(X, Y, field)
    # pl.show()



#if __name__ == '__main__':
#    write_gals(0, 'Lo', 1.0)
